import os
from dotenv import load_dotenv, find_dotenv
_ = load_dotenv(find_dotenv())
# 智谱AI
api_key = os.environ.get('ZHIPUAI_API_KEY') 
# OpenAI
# api_key = os.environ.get('OPENAI_API_KEY')
if api_key is None:
    raise ValueError("API Key is not set in the .env file")

# from openai import OpenAI
# client = OpenAI(api_key=api_key)
from zhipuai import ZhipuAI
client = ZhipuAI(api_key=api_key)

def get_completion(prompt, model="glm-4-plus", temperature=0.01):
    messages = [{"role": "user", "content": prompt}]
    response = client.chat.completions.create(
        model=model,
        messages=messages,
        temperature=temperature  
    )
    return response.choices[0].message.content
  
while True:  
    user_input = input("请输入提示文字（输入'exit'退出程序）: ")  
    if user_input is None or user_input == "" :
        continue
    elif user_input.lower() == 'exit':  
        print("程序退出。")  
        break  
    else :
        res = get_completion(user_input, temperature=0.9)  
        print("\n===大语言模型completion：===") 
        print(res)
        print("============END============\n")  